#!/usr/bin/env python3
"""
Weld Data MCP 服务器 (MySQL版)
"""

import asyncio
import json
import logging
from typing import Any, Dict, List, Optional


import aiomysql
from mcp.server import Server
from mcp.server.stdio import stdio_server
from mcp.types import (
    # CallToolRequest,
    # CallToolResult,
    # ListToolsRequest,
    # ListToolsResult,
    Tool,
    TextContent,
)

# from tavily import TavilyClient

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Tavily API配置
# TAVILY_API_KEY = "Your API Key"


# ==== MySQL 配置（改成你的实际数据库信息） ====

MYSQL_CONFIG = {
    "host": "39.98.163.114",
    "port": 8891,
    "user": "root",
    "password": "zedx1234",
    "db": "zedx_prodb",
}

class WeldMCPServer:
    """焊口查询 MCP Server"""

    def __init__(self):
        self.server = Server("weld-data")
        self.setup_handlers()

    def setup_handlers(self):
        """设置 MCP 处理器"""

        @self.server.list_tools()
        async def list_tools() -> List[Tool]:
            """列出可用的工具"""
            return [
                Tool(
                    name="query_weld_data",
                    description="根据人员姓名和日期，查询该人员当天的焊口编号及数量",
                    inputSchema={
                        "type": "object",
                        "properties": {
                            "person_name": {
                                "type": "string",
                                "description": "人员姓名"
                            },
                            "date": {
                                "type": "string",
                                "description": "日期,格式YYYY-MM-DD"
                            }
                        },
                        "required": ["person_name", "date"]
                    }
                )
            ]

        @self.server.call_tool()
        async def call_tool(name: str, arguments: Dict[str, Any]) -> List[TextContent]:
            """调用工具"""
            try:
                if name == "query_weld_data":
                    return await self._query_weld_data(arguments)
                else:
                    raise ValueError(f"Unknown tool: {name}")
            except Exception as e:
                logger.error(f"Tool call failed: {e}")
                return [TextContent(type="text", text=f"Error: {str(e)}")]

    async def _query_weld_data(self, arguments: Dict[str, Any]) -> List[TextContent]:
        """查询焊口数据"""
        person_name = arguments.get("person_name")
        date = arguments.get("date")

        if not person_name or not date:
            return [TextContent(type="text", text="Error: person_name 和 date 是必填项")]

        logger.info(f"查询 {person_name} 在 {date} 的焊口数据")

        try:
            async with aiomysql.connect(**MYSQL_CONFIG) as conn:
                async with conn.cursor() as cur:
                    # 查人员ID
                    await cur.execute("SELECT id FROM persons WHERE name=%s", (person_name,))
                    row = await cur.fetchone()
                    if not row:
                        return [TextContent(type="text", text=f"未找到人员 {person_name}")]
                    person_id = row[0]

                    # 查焊口
                    await cur.execute(
                        "SELECT weld_id FROM welds WHERE person_id=%s AND weld_date=%s",
                        (person_id, date),
                    )
                    weld_rows = await cur.fetchall()
                    weld_ids = [r[0] for r in weld_rows]

                    result = {
                        "person_name": person_name,
                        "person_id": person_id,
                        "date": date,
                        "count": len(weld_ids),
                        "weld_ids": weld_ids,
                    }

                    return [TextContent(type="text", text=json.dumps(result, ensure_ascii=False, indent=2))]

        except Exception as e:
            logger.error(f"MySQL query failed: {e}")
            return [TextContent(type="text", text=f"MySQL query failed: {str(e)}")]

    async def run(self):
        """运行 MCP 服务器"""
        logger.info("Starting Weld MCP Server (MySQL)...")
        async with stdio_server() as (read_stream, write_stream):
            await self.server.run(
                read_stream,
                write_stream,
                self.server.create_initialization_options()
            )


async def main():
    """主函数"""
    server = WeldMCPServer()
    await server.run()


if __name__ == "__main__":
    asyncio.run(main())


    
#     def _format_search_result(self, response: Dict[str, Any], query: str) -> Dict[str, Any]:
#         """格式化搜索结果"""
#         try:
#             results = response.get("results", [])
#             answer = response.get("answer", "")
            
#             # 格式化结果列表
#             formatted_results = []
#             for result in results:
#                 formatted_results.append({
#                     "title": result.get("title", ""),
#                     "url": result.get("url", ""),
#                     "content": result.get("content", ""),
#                     "score": result.get("score", 0)
#                 })
            
#             # 生成摘要答案
#             if answer:
#                 summary = answer
#             elif formatted_results:
#                 # 如果没有直接答案，从结果中生成摘要
#                 summary = self._generate_summary(formatted_results, query)
#             else:
#                 summary = f"No relevant information found for '{query}'."
            
#             return {
#                 "query": query,
#                 "answer": summary,
#                 "results": formatted_results,
#                 "total_results": len(formatted_results)
#             }
            
#         except Exception as e:
#             logger.error(f"Failed to format search result: {e}")
#             return {
#                 "query": query,
#                 "answer": f"Result processing failed: {str(e)}",
#                 "results": [],
#                 "error": str(e)
#             }
    
#     def _generate_summary(self, results: List[Dict[str, Any]], query: str) -> str:
#         """从搜索结果生成摘要"""
#         if not results:
#             return f"No relevant information found for '{query}'."
        
#         # 取前3个结果的内容片段
#         summary_parts = []
#         for i, result in enumerate(results[:3]):
#             content = result.get("content", "").strip()
#             if content:
#                 # 截取前200个字符
#                 snippet = content[:200] + "..." if len(content) > 200 else content
#                 summary_parts.append(f"{i+1}. {snippet}")
        
#         if summary_parts:
#             return f"Based on search results for '{query}':\n\n" + "\n\n".join(summary_parts)
#         else:
#             return f"Found search results for '{query}', but content summary is not available."
    
#     async def run(self):
#         """运行MCP服务器"""
#         logger.info("Starting Tavily MCP Server...")
#         async with stdio_server() as (read_stream, write_stream):
#             await self.server.run(
#                 read_stream,
#                 write_stream,
#                 self.server.create_initialization_options()
#             )

# async def main():
#     """主函数"""
#     server = TavilyMCPServer()
#     await server.run()

# if __name__ == "__main__":
#     asyncio.run(main())

